#include "mincut_segmenter.h"
#include <queue>
#include <xtd/top_selector.h>
//#include <xtd/group.h>
#include <imglib/imageclass.h>
#include <imglib/bmpmorph.h>
//#include <boost/graph/edmunds_karp_max_flow.hpp>
#include <boost/graph/push_relabel_max_flow.hpp>

namespace trm {

namespace mincut {

using std::clog;
using std::endl;
using std::flush;

segmenter::segmenter(mswin::timing_tree * timing, bool be_silent) :
	color(get(vertex_color, g)),
	cap(get(edge_capacity, g)),
	residual_capacity(get(edge_residual_capacity, g)),
	rev_edge(get(edge_reverse, g)),
	timing_tree(timing),
	silent(be_silent)
{
	if (!silent)
		clog << "MinCut segmenting" << endl;

	source_vertex = add_vertex(g);
	sink_vertex = add_vertex(g);
}

void segmenter::create_vmap(	
	mask const & obligatory_foreground_pixels,
	mask const & obligatory_background_pixels,
	vmap_type & vmap)
{
	mswin::timing_mark mark(timing_tree, "creating vmap");

	size_extent dim = obligatory_foreground_pixels.dim();
	assert (dim == obligatory_background_pixels.dim());
	vmap.resize(dim);

	for (vmap_type::iterator it = vmap.begin(); it != vmap.end(); ++it) {
		if (obligatory_foreground_pixels[it])
			*it = source_vertex;
		else if (obligatory_background_pixels[it])
			*it = sink_vertex;
		else
			*it = add_vertex(g);
	}
}

void segmenter::create_vmap(mask const & m, vmap_type & vmap)
{
	mswin::timing_mark mark(timing_tree, "creating vmap 1");

	vmap.resize(m.dim());

	for (vmap_type::iterator it = vmap.begin(); it != vmap.end(); ++it)
		*it = m[it] ? source_vertex : sink_vertex;
}

void segmenter::establish_links(
	byte_image const & center_image,
	byte_image const & periphery_image,
	vmap_type const & center_vmap,
	vmap_type const & periphery_vmap,
	size_t hw, size_t nn, float aspect
)
{
	mswin::timing_mark mark(timing_tree, "establishing links");
	bool one_image =
		center_image.data() == periphery_image.data();
	if (one_image)
		assert (center_vmap.data() == periphery_vmap.data());

	size_extent dim = center_image.dim();
	assert (dim == periphery_image.dim());
	assert (dim == center_vmap.dim());
	assert (dim == periphery_vmap.dim());

	mask periphery_fixed_pixels(dim);
	for (mask::iterator mit = periphery_fixed_pixels.begin(); mit != periphery_fixed_pixels.end(); ++mit) {
		vertex_descriptor u = periphery_vmap[mit];
		* mit = u == source_vertex || u == sink_vertex;
	}
	mask eroded_periphery_fixed_pixels(dim);
	img::erode8(eroded_periphery_fixed_pixels, periphery_fixed_pixels, hw);

	typedef float weight;

	//presort all possible links acording to their minimum weight
	typedef std::pair<weight, int_point> weight_link;
	std::vector<weight_link> weight_link_vector;
	weight_link_vector.reserve((2*hw+1) * (2*hw+1));
	int_point d;
	for (d.y = -(int)hw; d.y <= (int)hw; ++d.y)
		for (d.x = -(int)hw; d.x <= (int)hw; ++d.x)
			weight_link_vector.push_back(weight_link(aspect * (d.x*d.x + d.y*d.y), d));
	std::sort(weight_link_vector.begin(), weight_link_vector.end());

	typedef std::pair<weight, vertex_descriptor> weight_vertex;
	xtd::top_selector<weight_vertex> mins;
	size_t ec = 0;
	for (byte_image::const_iterator cit = center_image.begin(); cit != center_image.end(); ++cit) {
		int_point cp = size_point(cit);
		vertex_descriptor ci = center_vmap[cp];
		bool source_center = ci == source_vertex;
		bool sink_center = ci == sink_vertex;
		bool fixed_center = source_center || sink_center;
		if (fixed_center && eroded_periphery_fixed_pixels[cit])
			continue; //all possible edges are from source/sink to source/sink
		mins.start(nn);
		typedef img::BGRA<int> int_color;
		int_color cc = *cit;
		int_rect frame(0, 0, dim.sx, dim.sy);
		for (uint i = 0; i < weight_link_vector.size(); ++i) {
			weight_link const & wl = weight_link_vector[i];
			if (one_image && wl.second == int_point(0,0))
				continue; //otherwise it will definitely (weight is zero) be a link from center vertex to itself
			if (mins.full() && mins.worst_selected_value().first < wl.first)
				continue;
			int_point dp = int_point(cp) + wl.second;
			if (!frame.inside(dp))
				continue;
			int_color dc = int_color(periphery_image[dp]) - cc;
			mins.push_back(weight_vertex(wl.first + dc.sqr(), periphery_vmap[dp]));
		}
		mins.finish();
		// create edges to first n pixels with lowest weights
		size_t n = std::min(nn, mins.size());
		if (source_center) {
			//create only out edges, and only to the ordinary pixels
			for (size_t i = 0; i < n; ++i) {
				vertex_descriptor opponent = mins[i].second;
				if (opponent != sink_vertex && opponent != source_vertex) {
					inc_edge(source_vertex, opponent, n - i);
					++ec;
				}
			}
		}
		else if (sink_center) {
			//create only in edges, and only from the ordinary pixels
			for (size_t i = 0; i < n; ++i) {
				vertex_descriptor opponent = mins[i].second;
				if (opponent != sink_vertex && opponent != source_vertex) {
					inc_edge(opponent, sink_vertex, n - i);
					++ec;
				}
			}
		} else {
			//create bidirectional edges to ordinary pixels and single directional edges to source/sink
			for (size_t i = 0; i < n; ++i) {
				vertex_descriptor opponent = mins[i].second;
				if (opponent == source_vertex)
					inc_edge(source_vertex, ci, n - i);
				else if (opponent == sink_vertex)
					inc_edge(ci, sink_vertex, n - i);
				else
					inc_bidir_edge(ci, opponent, n - i);
				++ec;
			}
		}
	}
	if (!silent)
		clog << "edges added " << ec << endl;
}

/*
enum neighbour_colors {
	no_colored_neighbours = 0,
	has_white_neighbour = 1,
	has_black_neighbour = 2
};

struct vset : public xtd::patch<>
{
	int nc; //neighbour_colors
	bool visited; ///< property of individual vertex
	vset() : nc(no_colored_neighbours), visited(false) {}
	void subdue(vset & child)
		{ nc |= child.nc; }
};
*/

void segmenter::solve()
{
	mswin::timing_mark mark(timing_tree, "finding max-flow & min-cut");

	if (!silent)
		clog << "determining max flow..." << flush;
	//capacity flow = edmunds_karp_max_flow(g, source_vertex, sink_vertex);
	capacity flow = push_relabel_max_flow(g, source_vertex, sink_vertex);
	if (!silent)
		clog << "done\n"
			"max flow = " << flow << endl;

	size_t numv = num_vertices();
	// mark all vertices as gray (ambiguous)
	for (vertex_descriptor i = 0; i < numv; ++i)
		color[i] = gray_color;

	// mark pixels connected to the source with white color
	// (breadth-first search along not-saturated out-edges)
	std::queue<vertex_descriptor> Q;
	color[source_vertex] = white_color;
	Q.push(source_vertex);
	while (!Q.empty()) {
		vertex_descriptor u = Q.front(); Q.pop();
		Graph::out_edge_iterator ei, ei_end;
		for (tie(ei, ei_end) = out_edges(u, g); ei != ei_end; ++ei) {
			if (residual_capacity[*ei] == 0)
				continue;
			vertex_descriptor v = target(*ei, g);
			if (color[v] != gray_color)
				continue;
			color[v] = white_color;
			Q.push(v);
		}
	}

	// ... and to the sink with black color
	// (breadth-first search along not-saturated in-edges)
	color[sink_vertex] = black_color;
	Q.push(sink_vertex);
	while (!Q.empty()) {
		vertex_descriptor u = Q.front(); Q.pop();
		Graph::out_edge_iterator ei, ei_end;
		for (tie(ei, ei_end) = out_edges(u, g); ei != ei_end; ++ei) {
			if (residual_capacity[rev_edge[*ei]] == 0)
				continue;
			vertex_descriptor v = target(*ei, g);
			if (color[v] != gray_color)
				continue;
			color[v] = black_color;
			Q.push(v);
		}
	}

/*		experiments shown that the following code does not improve result

	// group remaining gray vertices into disjoint sets along 
	// all edges of the graph (including saturated);
	// for each set of vertices find colors of neighbours;
	typedef std::vector<vset> vset_vector;
	vset_vector vsets(numv);
	// (breadth-first search along all out-edges)
	vsets[source_vertex].visited = true;
	Q.push(source_vertex);
	while (!Q.empty()) {
		vertex_descriptor u = Q.front(); Q.pop();
		Graph::out_edge_iterator ei, ei_end;
		for (tie(ei, ei_end) = out_edges(u, g); ei != ei_end; ++ei) {
			if (cap[*ei] == 0)
				continue; //duplicate edge
			vertex_descriptor v = target(*ei, g);
			default_color_type cu = color[u], cv = color[v];
			if (cu == gray_color) {
				switch (cv) {
				default:
					assert(false);
				case gray_color:
					xtd::union_sets(vsets, u, v);
					break;
				case white_color:
					xtd::find_root(vsets, u).nc |= has_white_neighbour;
					break;
				case black_color:
					xtd::find_root(vsets, u).nc |= has_black_neighbour;
					break;
				}
			}
			else if (cv == gray_color) {
				switch (cu) {
				default:
					assert(false);
				case white_color:
					xtd::find_root(vsets, v).nc |= has_white_neighbour;
					break;
				case black_color:
					xtd::find_root(vsets, v).nc |= has_black_neighbour;
					break;
				}
			}
			if (vsets[v].visited)
				continue;
			vsets[v].visited = true;
			Q.push(v);
		}
	}

	// mark all gray vertices, which have only neighbours of one color with the neighbour's color
	size_t recolored = 0;
	for (vertex_descriptor u = 0; u < numv; ++u) {
		if (color[u] != gray_color)
			continue;
		int nc = xtd::find_root(vsets, u).nc;
		if (nc == has_white_neighbour) {
			color[u] = white_color;
			++recolored;
		}
		else if (nc == has_black_neighbour) {
			color[u] = black_color;
			++recolored;
		}
	}
	if (!p.silent)
		clog << "recolored pixels: " << recolored << endl;*/
}

void segmenter::get_frame_mask(vmap_type const & vmap, default_color_type col, mask & m) const
{
	if (m.dim() != vmap.dim())
		m.resize(vmap.dim());
	for (mask::iterator it = m.begin(); it != m.end(); ++it)
		*it = color(vmap[it]) == col;
}

void segmenter::inc_frame_gray_mask(vmap_type const & vmap, gray_mask & m) const
{
	assert (m.dim() == vmap.dim());
	for (gray_mask::iterator it = m.begin(); it != m.end(); ++it) {
		default_color_type col = color(vmap[it]);
		unsigned inc = (col == white_color) ? 2 : ((col == gray_color) ? 1 : 0);
		*it = 2 * *it + inc;
	}
}

void segmenter::double_capacity_of_saturated_edges()
{
	Graph::edge_iterator ei, ei_end;
	for (tie(ei, ei_end) = edges(g); ei != ei_end; ++ei) {
		if (cap[*ei] > 0 && residual_capacity[*ei] == 0)
			cap[*ei] *= 2;
	}
}

} //mincut

} //param
